/**
 * @file bernoulli_distr.h
 *
 * The class to represent a Bernoulli distribution
 * 
 * @author Dahua Lin
 */

#ifndef SLIPP_BERNOULLI_DISTR_H_
#define SLIPP_BERNOULLI_DISTR_H_

#include <slipp/prob_base/pdistr_base.h>

namespace sli
{

	template<class RStream>
	class bernoulli_sampler
	{
	public:
		SLIPP_DEFINE_SCALAR_DISTR_TYPES(bool)

	public:
		bernoulli_sampler(RStream& rstream, double p)
		: m_rng(rstream, p)
		{
		}

		sample_type draw()
		{
			return m_rng.get();
		}

		sample_collection_type draw(size_t n)
		{
			sample_collection_type v((index_t)n);
			m_rng.get_m(n, v.data());
			return v;
		}

	private:
		bernoulli_rng<RStream> m_rng;
	};


	class bernoulli_distr
	{
	public:
		SLIPP_DEFINE_SCALAR_DISTR_TYPES(bool)

	public:
		class param_type
		{
		public:
			param_type() : m_p(0.5) { }
			param_type(const double& p) : m_p(p) { }

			double p() const { return m_p; }
			double q() const { return 1.0 - m_p; }

			bool operator == (const param_type& rhs) const { return m_p == rhs.m_p; }
			bool operator != (const param_type& rhs) const { return !operator == (rhs); }

		private:
			double m_p;
		};

	public:
		// constructors

		bernoulli_distr()
		: m_param()
		{
		}

		bernoulli_distr(const double& p)
		: m_param(p)
		{
		}

		bernoulli_distr(const param_type& param)
		: m_param(param)
		{
		}

		// parameters

		const param_type& param() const
		{
			return m_param;
		}

		double p() const
		{
			return m_param.p();
		}

		double q() const
		{
			return m_param.q();
		}

		// statistics

		double mean() const
		{
			return p();
		}

		double var() const
		{
			return p() * q();
		}

		double entropy() const
		{
			return p() > 0 && p() < 1 ? - (p() * math::log(p()) + q() * math::log(q())) : 0.0;
		}

	public:
		class pdf_evaluator_type
		{
		public:
			typedef bernoulli_distr::sample_type sample_type;
			typedef bernoulli_distr::sample_collection_type sample_collection_type;

			pdf_evaluator_type(const param_type& param)
			: m_p(param.p()), m_q(param.q())
			{
			}

			SLIPP_ENSURE_INLINE double logpdf(const scalar_type& x) const
			{
				return x ? math::log(m_p) : math::log(m_q);
			}

			SLIPP_ENSURE_INLINE double pdf(const scalar_type& x) const
			{
				return x ? m_p : m_q;
			}

			darow_t pdf(const sample_collection_type& X) const
			{
				return map(_selector(m_p, m_q), X);
			}

			darow_t logpdf(const sample_collection_type& X) const
			{
				return map(_selector(math::log(m_p), math::log(m_q)), X);
			}

		private:
			struct _selector
			{
				typedef double result_type;
				double v1, v0;

				_selector(double v1_, double v0_) : v1(v1_), v0(v0_) { }

				double operator() (bool b) const
				{
					return b ? v1 : v0;
				}
			};

		private:
			double m_p;
			double m_q;
		};

		pdf_evaluator_type get_pdf_evaluator() const
		{
			return pdf_evaluator_type(m_param);
		}

	public:
		// sampling

		template<class RStream>
		bernoulli_sampler<RStream> get_sampler(RStream& rstream) const
		{
			return bernoulli_sampler<RStream>(rstream, p());
		}

	private:
		param_type m_param;

	}; // end class bernoulli_distr


	// traits

	SLIPP_DEFINE_BASIC_DISTR_TRAITS(bernoulli_distr, bernoulli_sampler)

}

#endif 

